import cv2
import torch
import numpy as np
from PIL import Image
from torchvision.transforms import ToTensor


def process_video(input_path):
    """处理视频文件并返回输出视频路径"""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = torch.load('models/model.pkl').to(device)
    model.eval()
    transform = ToTensor()

    cap = cv2.VideoCapture(input_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    output_path = input_path.replace('.', '_output.')
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

    labels = {
        0: "拳头",
        1: "手掌",
        2: "剪刀手",
        3: "OK手势",
        4: "竖起大拇指"
    }

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        # 处理帧
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        pil_img = Image.fromarray(rgb_frame)
        img_tensor = transform(pil_img).unsqueeze(0).to(device)

        with torch.no_grad():
            output = model(img_tensor)
            _, predicted = torch.max(output.data, 1)
            label = labels.get(predicted.item(), "未知手势")

        # 添加识别结果到帧
        cv2.putText(frame, f"手势: {label}", (20, 50),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        out.write(frame)

    cap.release()
    out.release()
    return output_path